import json
import os
from typing import Dict, Any
from request_map_generator import MapGenerator3D, build_map_from_environmental_info
from multi_uav_planner_3D_waypoints import MultiUAVPlanner  
import pickle
from typing import List, Dict, Tuple, Optional
import csv


def load_map_data(filename):
    """加载保存的地图数据"""
    if os.path.exists(filename):  # 检查文件是否存在
        with open(filename, 'rb') as f:
            map_data = pickle.load(f)
        print(f"已加载地图数据: {filename}")
        print(f"地图尺寸: {map_data['x_width']}x{map_data['y_width']}x{map_data['z_width']}")
        return map_data
    else:
        return None  # 如果没有找到文件，返回 None


class MissionExecutor:
    def __init__(self, 
                 map_resolution: float = 20.0, 
                 robot_radius: float = 10.0, 
                 map_file: str = "test_map_from_env_info.pkl",
                 # 区域覆盖规划器参数（传递给MultiUAVPlanner）----------------
                 patrol_speed: float = 10.0,
                 min_patrol_speed: float = 5.0,
                 max_patrol_speed: float = 25.0,
                 target_patrol_time=600.0,
                 time_buffer: float = 120.0,
                 default_altitude: float = 10.0,
                 sweep_width: float = 12.0,
                 max_segment_length: float = 20.0,
                 path_boundary_standoff: float = 5.0,
                 # 3D规划器参数（传递给MultiUAVPlanner）----------------
                 speed: float = 20.0,
                 safe_dis: float = 2.0,
                 min_planning_speed: float = 5.0,
                 max_planning_speed: float = 25.0,
                 speed_increase_factor: float = 1.1,
                 max_t: int = 3000,
                 goal_tolerance: float = 2.0,
                 heuristic_weight: float = 1.5):
        """
        任务执行器初始化，支持传递MultiUAVPlanner所需的所有关键参数
        
        Args:
            map_resolution: 地图分辨率（米/栅格）
            robot_radius: 无人机半径（用于避障计算）
            map_file: 预加载/保存的地图文件路径
            # 区域覆盖规划器参数（巡航、扫查任务用）
            patrol_speed: 巡航初始速度（m/s）
            min_patrol_speed: 巡航速度下限（m/s）
            max_patrol_speed: 巡航速度上限（m/s）
            target_patrol_time : float, 默认600.0巡航任务时间约束（任务必须在此时间内完成），单位：秒（s）
            time_buffer: 任务时间缓冲（s）
            default_altitude: 默认飞行高度（m）
            sweep_width: 传感器扫描宽度（m）
            max_segment_length: 巡航路径段最大长度（m）
            path_boundary_standoff: 巡航路径与边界安全距离（m）
            # 3D规划器参数（路径搜索、避障用）
            speed: 3D规划初始速度（m/s）
            safe_dis: 与障碍物安全距离（m）
            min_planning_speed: 规划速度下限（m/s）
            max_planning_speed: 规划速度上限（m/s）
            speed_increase_factor: 速度调整系数（>1）
            max_t: 最大规划时间（s，超时判定失败）
            goal_tolerance: 目标点位置容差（m）
            heuristic_weight: A*启发式权重（1.1~2.0，值越大搜索越快）
        """
        # 基础参数
        self.map_resolution = map_resolution
        self.robot_radius = robot_radius
        self.map_file = map_file

        # 区域覆盖规划器参数
        self.patrol_speed = patrol_speed
        self.min_patrol_speed = min_patrol_speed
        self.max_patrol_speed = max_patrol_speed
        self.time_buffer = time_buffer
        self.target_patrol_time = target_patrol_time
        self.default_altitude = default_altitude
        self.sweep_width = sweep_width
        self.max_segment_length = max_segment_length
        self.path_boundary_standoff = path_boundary_standoff
        # 3D规划器参数
        self.speed = speed
        self.safe_dis = safe_dis
        self.min_planning_speed = min_planning_speed
        self.max_planning_speed = max_planning_speed
        self.speed_increase_factor = speed_increase_factor
        self.max_t = max_t
        self.goal_tolerance = goal_tolerance
        self.heuristic_weight = heuristic_weight
        
        # 加载地图并初始化规划器
        self.map_env = load_map_data(self.map_file)  # 尝试加载地图数据
        self.ref_lat = self.map_env['ref_lat'] if self.map_env else None
        self.ref_lon = self.map_env['ref_lon'] if self.map_env else None
        # 如果地图数据已加载，初始化规划器
        self.planner = self._init_planner() if self.map_env else None

    def _init_planner(self) -> MultiUAVPlanner:
        """内部工具方法：初始化MultiUAVPlanner"""
        if not self.map_env or self.ref_lat is None or self.ref_lon is None:
            raise ValueError("地图数据或参考经纬度缺失，无法初始化规划器")
        # 传递所有关键参数给MultiUAVPlanner
        return MultiUAVPlanner(
            map_env=self.map_env,
            ref_lat=self.ref_lat,
            ref_lon=self.ref_lon,
            # 区域覆盖规划器参数
            patrol_speed=self.patrol_speed,
            min_patrol_speed=self.min_patrol_speed,
            max_patrol_speed=self.max_patrol_speed,
            time_buffer=self.time_buffer,
            target_patrol_time=self.target_patrol_time,
            default_altitude=self.default_altitude,
            sweep_width=self.sweep_width,
            max_segment_length=self.max_segment_length,
            path_boundary_standoff=self.path_boundary_standoff,
            # 3D规划器参数
            speed=self.speed,
            safe_dis=self.safe_dis,
            min_planning_speed=self.min_planning_speed,
            max_planning_speed=self.max_planning_speed,
            speed_increase_factor=self.speed_increase_factor,
            max_t=self.max_t,
            goal_tolerance=self.goal_tolerance,
            heuristic_weight=self.heuristic_weight
        )

      # ---------------- 外部 JSON → 内部任务 ----------------
    def adapt_external_json(self, external_json: Dict[str, Any]):
        """解析 template.json 风格的 `routeToBePlannedInfos`。

        返回 (tasks, merged_environmental_info, bounds) 三元组。
        - tasks: 与 adapt_external_json 返回相同格式的任务列表
        - merged_environmental_info: 合并后的 environmentalInfo 字段
        - bounds: dict 包含 min_lon,max_lon,min_lat,max_lat 键，表示所有任务的经纬度边界
        """
        tasks = []
        device_model_map = external_json.get("deviceModelParametersMap", {}) or {}

        # 准备合并的 environment info
        merged_env = {"obstacles": [], "waterways": [], "roads": [], "noFlyZones": [], "sensitiveAreas": []}

        lons = []
        lats = []
        alts = []

        for entry in external_json.get("routeToBePlannedInfos", []):
            # 1) 只处理旋翼机
            if entry.get("deviceType") != "multi_rotor":
                continue

            # 合并条目的 environment info（若存在）。
            # 注意：将 noFlyZones 与 sensitiveAreas 也当作障碍物处理，以便建图时统一考虑
            ei = entry.get("environmentalInfo") or {}
            if ei:
                if ei.get("obstacles"):
                    merged_env.setdefault("obstacles", []).extend(ei.get("obstacles"))
                if ei.get("waterways"):
                    merged_env.setdefault("waterways", []).extend(ei.get("waterways"))
                if ei.get("roads"):
                    merged_env.setdefault("roads", []).extend(ei.get("roads"))
                # 把 noFlyZones 与 sensitiveAreas 也加入 obstacles 列表
                if ei.get("noFlyZones"):
                    merged_env.setdefault("noFlyZones", []).extend(ei.get("noFlyZones"))
                    merged_env.setdefault("obstacles", []).extend(ei.get("noFlyZones"))
                if ei.get("sensitiveAreas"):
                    merged_env.setdefault("sensitiveAreas", []).extend(ei.get("sensitiveAreas"))
                    merged_env.setdefault("obstacles", []).extend(ei.get("sensitiveAreas"))

            # 设备信息统一为 deviceInfos 列表
            device_infos = entry.get("deviceInfos") or ([{"deviceId": entry.get("deviceId")}] if entry.get("deviceId") else [])
            if not device_infos:
                continue
            uav_id = device_infos[0].get("deviceId")

            device_model = entry.get("deviceModel")

            # 统一途径点字段名（waypoints / wayPoints）
            wps = entry.get("waypoints") or entry.get("wayPoints") or []

            # 收集起终点与途径点用于建图范围
            sp = entry.get("startPoint")
            ep = entry.get("endPoint")
            pts = []
            if sp:
                pts.append(sp)
                lons.append(sp.get("longitude"))
                lats.append(sp.get("latitude"))
                alts.append(sp.get("altitude", 0))
            for wp in wps:
                pts.append(wp)
                lons.append(wp.get("longitude"))
                lats.append(wp.get("latitude"))
                alts.append(wp.get("altitude", 0))
            if ep:
                pts.append(ep)
                lons.append(ep.get("longitude"))
                lats.append(ep.get("latitude"))
                alts.append(ep.get("altitude", 0))

            # 读取时间约束：优先使用条目 extraInfo.maxFlightTime，否则尝试从 deviceModelParametersMap 中取
            extra = entry.get("extraInfo") or {}
            mft = None
            if "maxFlightTime" in extra and extra["maxFlightTime"] is not None:
                # 假定外部若给出就是以秒为单位
                try:
                    mft = float(extra["maxFlightTime"])
                except Exception:
                    mft = None

            if mft is None and device_model:
                dm = device_model_map.get(device_model, {})
                if dm and dm.get("maxFlightTime") is not None:
                    # deviceModelParametersMap 中的 maxFlightTime 单位为分钟，转换为秒
                    try:
                        mft = float(dm.get("maxFlightTime")) * 60.0
                    except Exception:
                        mft = None

            # 最终兜底为 executor 的 target_patrol_time
            if mft is None:
                mft = float(self.target_patrol_time)

            # 读取设备型号的速度上限（vMaxHorizon）用于建议 patrol_speed / planning speed 上限
            v_max_horizon = None
            if device_model:
                dm = device_model_map.get(device_model, {})
                if dm and dm.get("vMaxHorizon") is not None:
                    try:
                        v_max_horizon = float(dm.get("vMaxHorizon"))
                    except Exception:
                        v_max_horizon = None

            # 根据 routeType 构造任务
            route_type = (entry.get("routeType") or "").lower()
            if route_type == "point":
                waypoint_sequence = []
                if sp:
                    waypoint_sequence.append({"lon": sp["longitude"], "lat": sp["latitude"], "alt": sp.get("altitude", 0)})
                for wp in wps:
                    waypoint_sequence.append({"lon": wp["longitude"], "lat": wp["latitude"], "alt": wp.get("altitude", 0)})
                if ep:
                    waypoint_sequence.append({"lon": ep["longitude"], "lat": ep["latitude"], "alt": ep.get("altitude", 0)})

                if len(waypoint_sequence) > 2:
                    task = {
                        "priority": entry.get("priority", 1),
                        "taskType": "multi_waypoint",
                        "uav_ids": [uav_id],
                        "time_constraint": int(mft),
                        "start": [{"uav_id": uav_id, "position": waypoint_sequence[0]}],
                        "waypoint_sequence": waypoint_sequence,
                        "deviceModel": device_model,
                        "deviceType": entry.get("deviceType"),
                        "suggested_max_speed": v_max_horizon
                    }
                else:
                    # 起点+终点
                    task = {
                        "priority": entry.get("priority", 1),
                        "taskType": "waypoint",
                        "uav_ids": [uav_id],
                        "time_constraint": int(mft),
                        "start": [{"uav_id": uav_id, "position": waypoint_sequence[0]}],
                        "goals": [{"uav_id": uav_id, "position": waypoint_sequence[1]}],
                        "deviceModel": device_model,
                        "deviceType": entry.get("deviceType"),
                        "suggested_max_speed": v_max_horizon
                    }

                tasks.append(task)

            elif route_type in ("area", "patrol"):
                # area 类型：把 entry 的 waypoints（或 wayPoints）当作巡航区域的顶点。
                vertices = wps or entry.get("routePoints") or entry.get("vertices") or []
                # 如果 vertices 是字典列表，按字段转换；否则保持空
                verts_out = []
                for v in vertices:
                    if isinstance(v, dict):
                        verts_out.append({"lon": v.get("longitude"), "lat": v.get("latitude"), "alt": v.get("altitude", 0)})
                task = {
                    "priority": entry.get("priority", 1),
                    "taskType": "patrol",
                    "uav_ids": [uav_id],
                    "time_constraint": int(mft),
                    "start": [{"uav_id": uav_id, "position": {"lon": sp.get("longitude"), "lat": sp.get("latitude"), "alt": sp.get("altitude", 0)} }] if sp else [],
                    "vertices": verts_out,
                    "deviceModel": device_model,
                    "deviceType": entry.get("deviceType"),
                    "suggested_max_speed": v_max_horizon
                }
                tasks.append(task)

        bounds = None
        if lons and lats:
            bounds = {"min_lon": min(filter(None, lons)), "max_lon": max(filter(None, lons)), "min_lat": min(filter(None, lats)), "max_lat": max(filter(None, lats))}

        # 合并顶层 environmentalInfo（如果有）
        top_env = external_json.get("environmentalInfo") or {}
        for k in ("obstacles", "noFlyZones", "sensitiveAreas"):
            if top_env.get(k):
                merged_env.setdefault(k, []).extend(top_env.get(k))

        return tasks, merged_env, bounds

    # ---------------- 整体执行流程 ----------------
    def execute(self, external_json: Dict[str, Any], map_file: str = "mission_map.pkl", show_visualization: bool = True):
        """
        执行任务规划
        
        Args:
            external_json: 外部任务JSON
            map_file: 地图文件名
            show_visualization: 是否显示可视化图表
            
        Returns:
            tuple: (gps_trajectories, world_paths) - GPS轨迹和世界坐标路径
        """
        # 1. 先把外部 JSON 转为内部任务与环境信息（tasks, merged_env, bounds）
        tasks, merged_env, bounds = self.adapt_external_json(external_json)
        if not tasks:
            raise ValueError("外部 JSON 中没有可用任务")

        # 2. 为每个任务分配对应的速度上限：优先使用任务的 suggested_max_speed（来自 deviceModelParametersMap），
        #    否则使用执行器全局上限 self.max_planning_speed / self.max_patrol_speed。
        for t in tasks:
            suggested = t.get("suggested_max_speed")
            # patrol 任务使用巡航速度限制；point/waypoint 使用规划速度限制
            if t.get("taskType") == "patrol":
                if suggested:
                    try:
                        # 使用设备建议上限
                        t["max_patrol_speed"] = float(suggested)
                    except Exception:
                        t["max_patrol_speed"] = float(self.max_patrol_speed)
                else:
                    t["max_patrol_speed"] = float(self.max_patrol_speed)
            else:
                # 对于 waypoint / multi_waypoint 等点任务，限制规划速度
                if suggested:
                    try:
                        # 使用设备建议上限
                        t["max_planning_speed"] = float(suggested)
                    except Exception:
                        t["max_planning_speed"] = float(self.max_planning_speed)
                else:
                    t["max_planning_speed"] = float(self.max_planning_speed)

        # 3. 如果已有地图，尝试复用并确保 planner 已初始化；否则根据 merged_env 和 bounds 建图
        if self.map_env:
            print("已存在地图数据，跳过建图。")
            if not self.planner:
                try:
                    self.planner = self._init_planner()
                except Exception as e:
                    raise RuntimeError(f"地图已存在但初始化规划器失败: {e}")
        else:
            if not bounds:
                raise ValueError("缺少 mission 边界信息，无法建图")
            print("正在根据任务集合生成地图...")
            ref_lat = 0.5 * (bounds["min_lat"] + bounds["max_lat"])
            ref_lon = 0.5 * (bounds["min_lon"] + bounds["max_lon"])

            generated_map = build_map_from_environmental_info(
                map_resolution=self.map_resolution,
                robot_radius=self.robot_radius,
                output_filename=map_file,
                env_info=merged_env,
                mission_bounds=bounds,
                ref_lat=ref_lat,
                ref_lon=ref_lon,
            )

            # 更新已有的地图数据并初始化 planner
            self.map_env = load_map_data(map_file)
            self.ref_lat = self.map_env['ref_lat']
            self.ref_lon = self.map_env['ref_lon']
            self.planner = self._init_planner()

        # 4. 把内部任务和地图传给多机规划器
        print("调用多机规划算法...")
        gps_trajectories, world_paths = self.planner.process_tasks(tasks)

        # ---------------- 给GPS轨迹添加建议速度 ----------------
        gps_with_speed = {}
        for uav_id in gps_trajectories.keys():
            if uav_id not in world_paths:
                print(f"无人机{uav_id}无世界路径，无法添加速度")
                gps_with_speed[uav_id] = gps_trajectories[uav_id]
                continue
            # 调用速度计算方法（传入 uav_id 以使用 per-UAV 上限）
            gps_with_speed[uav_id] = self.planner._add_speed_to_gps_trajectory(
                gps_traj=gps_trajectories[uav_id],
                world_path=world_paths[uav_id],
                uav_id=uav_id
            )
        # 更新为带速度的GPS轨迹
        gps_trajectories = gps_with_speed
        # 可视化
        if show_visualization and world_paths:
            print("\n开始可视化...")
            try:
                self.planner.plot_static_world_trajectories_3d(world_paths)
                self.planner.animate_world_trajectories_3d(world_paths, interval=100)
            except Exception as e:
                print(f"可视化过程中出现错误: {e}")

        return gps_trajectories, world_paths


if __name__ == "__main__":
    # 1. 加载外部任务JSON
    with open("template.json", "r", encoding="utf-8") as f:
        external_json = json.load(f)

    # 2. 初始化任务执行器（示例：自定义关键参数，适配特定场景）
    executor = MissionExecutor(
        map_resolution=20.0,          # 地图分辨率提高（更精细）
        robot_radius=12.0,             # 无人机半径减小
        map_file="test_map_from_env_info.pkl",
        # ---------------- 自定义区域覆盖规划器参数 ----------------
        patrol_speed=15.0,            # 巡航速度提高到15m/s
        target_patrol_time=900.0,    # 目标巡航时间延长到15分钟
        time_buffer=180.0,            # 时间缓冲延长到3分钟
        default_altitude=15.0,        # 默认飞行高度提高到15m
        max_segment_length=50.0,     # 巡航路径段最大长度
        sweep_width=40.0,             # 传感器扫描宽度扩大到40m
        path_boundary_standoff=10.0,   # 与边界安全距离扩大到10m
        # ---------------- 自定义3D规划器参数 ----------------
        speed=20.0,                   # 3D规划初始速度提高到20m/s
        safe_dis=20.0,                 # 与障碍物安全距离
        heuristic_weight=1.8,         # 启发式权重提高（优先搜索速度）
        max_t=3600                    # 最大规划时间延长到1小时
    )
    
    # 3. 执行规划并自动显示可视化
    gps_result, world_result = executor.execute(
        external_json, 
        map_file="test_map_from_env_info.pkl", 
        show_visualization=False  # 设置为False可禁用可视化
    )

    # 4. 输出规划结果到CSV
    print("\n规划结果:")
    if gps_result:
        # 创建输出目录（如果不存在）
        output_dir = "mission_results"
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # 为每个无人机生成CSV文件
        for uav_id, trajectory in gps_result.items():
            filename = os.path.join(output_dir, f"UAV_{uav_id}_trajectory.csv")
            
            with open(filename, 'w', newline='', encoding='utf-8') as csvfile:
                writer = csv.writer(csvfile)
                # 写入表头
                writer.writerow(['Latitude', 'Longitude', 'Altitude', 'Speed (m/s)'])
                # 写入航点数据
                for waypoint in trajectory:
                    writer.writerow([
                        waypoint['lat'],
                        waypoint['lon'],
                        waypoint['alt'],
                        waypoint.get('speed', 0.0)  # 若速度缺失，默认0
                    ])
            
            print(f"无人机 {uav_id} 的航点已保存到: {filename}")
            print(f"  共 {len(trajectory)} 个航点")